# -*- coding: utf-8 -*-
"""An agent class that implements the ReAct algorithm. The agent will reason
and act iteratively to solve problems. More details can be found in the paper
https://arxiv.org/abs/2210.03629.
"""
from typing import Any, Callable, Optional, Union, Sequence

from loguru import logger
from .flex_service_toolkit import FlexServiceToolkit
from agentscope.exception import ResponseParsingError, FunctionCallError
from agentscope.agents import AgentBase
from agentscope.memory.temporary_memory import TemporaryMemory
from agentscope.message import Msg
from agentscope.parsers import MarkdownJsonDictParser
from agentscope.service.service_toolkit import ServiceFunction
from agentscope.service import (
    ServiceToolkit,
    ServiceResponse,
    ServiceExecStatus,
)
import json

IF_PROMPT = """
You are an AI assistant to decide the condition is true or false.
Make your decide based on user's input and the context.
"""


class IfAgent(AgentBase):
    """
    IfAgent, check the condition and execute the true or false branch.
    """
    Instance = {}
    def __init__(
        self,
        name: str,
        model_config_name: str,
        service_toolkit: FlexServiceToolkit = None,
        sys_prompt: str = "You're a helpful assistant.",
        max_iters: int = 10,
        verbose: bool = True,
        isSpeak:bool = False,
        **kwargs: Any,
    ) -> None:
        """
        Args:
            name (`str`):
                The name of the agent.
            sys_prompt (`str`):
                The system prompt of the agent.
            model_config_name (`str`):
                The name of the model config, which is used to load model from
                configuration.
            service_toolkit (`ServiceToolkit`):
                A `ServiceToolkit` object that contains the tool functions.
            max_iters (`int`, defaults to `10`):
                The maximum number of iterations of each chapter.
            verbose (`bool`, defaults to `True`):
                Whether to print the detailed information during reasoning and
                acting steps. If `False`, only the content in speak field will
                be print out.
            chapter_chain (`Sequence[str]`):
                The chapter chain to be performed. If None, the agent will add a 
                init chapter to construct the chapter chain.
        """
        super().__init__(
            name=name,
            sys_prompt=sys_prompt,
            model_config_name=model_config_name,
        )
        self.model_config_name =  model_config_name
        # TODO: To compatible with the old version, which will be deprecated
        #  soon

        if service_toolkit is not None:
            # service_toolkit.add(record,name=name)
            # service_toolkit.add(next_chapter,name=name)
            # for tool in file_tools:
            #     service_toolkit.add(tool, root_path=root_path)
            pass

        self.service_toolkit = service_toolkit
        self.verbose = verbose
        self.max_iters = max_iters

        if not sys_prompt.endswith("\n"):
            sys_prompt = sys_prompt + "\n"

        self.sys_prompt = "\n".join(
            [
                sys_prompt.format(name=self.name),
                IF_PROMPT,
            ],
        )

        
        # This memory is used to store all useful memory with the chapter chain switching
        #self.memory_on_chain = TemporaryMemory()
        self.memory.add(Msg("system", self.sys_prompt, role="system"))
        
        # Initialize a parser object to formulate the response from the model
        content_hint = {
            "thought": "what you thought",
            "decide": "true or false"
        }
        required_keys = ["thought","decide"]
        if isSpeak: 
            content_hint["speak"] = "what you speak"
            required_keys.append("speak")
        if self.service_toolkit is not None:
            content_hint['function'] = service_toolkit.tools_calling_format
            required_keys.append("function")

        self.parser = MarkdownJsonDictParser(
                content_hint=content_hint,
            required_keys=required_keys,
            # Only print the speak field when verbose is False
            keys_to_content=True if self.verbose else "speak",
        )

    def reply(self, x: Optional[Union[Msg, Sequence[Msg]]] = None, ) -> Msg:
            self.memory.add(x)           
            for _ in range(self.max_iters):
                # Step 1: Thought
                if self.verbose:
                    self.speak(f" ITER {_+1}, STEP 1: REASONING ".center(70, "#"))

                # Prepare hint to remind model what the response format is
                # Won't be recorded in memory to save tokens
                if self.service_toolkit is not None:
                    tools_hint_msg = Msg(
                        "system",
                        self.service_toolkit.tools_instruction,
                        role="system",
                        echo=self.verbose,
                    )
                else : tools_hint_msg = Msg("system", "No Tool Available", role="system", echo=self.verbose)
                hint_msg = Msg(
                    "system",
                    self.parser.format_instruction+"\n",
                    role="system",
                    echo=self.verbose,
                )
                prompt = self.model.format(self.memory.get_memory(), 
                    tools_hint_msg,
                    hint_msg)
                if self.verbose:
                    self.speak(f"Prompt".center(70, "#"))
                    self.speak(str(prompt))
                
                try:
                    if self.verbose:
                        self.speak(f"API Trigger".center(70, "#"))
                    res = self.model(
                        prompt,
                        parse_func=self.parser.parse,
                        max_retries=2,
                        ######################
                    )
                    
                    if self.verbose:
                        self.speak(f"Result Parsed".center(70, "#"))
                        self.speak(str(res.parsed))
                        print(res.parsed)
                        
                    self.memory.add(
                        Msg(
                            self.name,
                            self.parser.to_memory(res.parsed),
                            "assistant",
                        ),
                    )

                    # Print out the response
                    msg_returned = Msg(
                        self.name,
                        self.parser.to_content(res.parsed),
                        "assistant",
                    )
                    if self.verbose:
                        self.speak(f"Speak".center(70, "#"))
                        print(msg_returned)
                    self.speak(msg_returned)
                    return msg_returned
                
                    
                    if self.verbose:
                        self.speak(f"Function".center(70, "#"))
                        self.speak(str(res.parsed["function"]))
                        print(res.parsed["function"])

                    arg_function = res.parsed["function"]
                    if (
                        isinstance(arg_function, str)
                        and arg_function in ["[]", ""]
                        or isinstance(arg_function, list)
                        and len(arg_function) == 0
                    ):
                        # Only the speak field is exposed to users or other agents
                        continue

                # Only catch the response parsing error and expose runtime
                # errors to developers for debugging
                except ResponseParsingError as e:
                    # Print out raw response from models for developers to debug
                    response_msg = Msg(self.name, e.raw_response, "assistant")
                    self.speak(response_msg)

                    # Re-correct by model itself
                    error_msg = Msg("system", str(e), "system")
                    self.speak(error_msg)

                    self.memory.add([response_msg, error_msg])

                

    def function_call(self,  parsed_function) -> Msg:
                if self.verbose:
                    self.speak(f"ACTING ".center(70, "#"))

                # Parse, check and execute the tool functions in service toolkit
                try:
                    execute_results = self.service_toolkit.parse_and_call_func(
                        parsed_function# res.parsed["function"],
                    )   

                    # Note: Observing the execution results and generate response
                    # are finished in the next reasoning step. We just put the
                    # execution results into memory, and wait for the next loop
                    # to generate response.

                    # Record execution results into memory as system message
                    msg_res = Msg("system", execute_results, "system")
                    self.speak(msg_res)
                    self.chapter_memory.add(msg_res)

                except FunctionCallError as e:
                    # Catch the function calling error that can be handled by
                    # the model
                    error_msg = Msg("system", str(e), "system")
                    self.speak(error_msg)
                    self.chapter_memory.add(error_msg)

                except Exception as e:
                    # Catch the function calling error that can not be handled
                    error_msg = Msg("system", str(e), "system")
                    self.speak(error_msg)
                    self.chapter_memory.add(error_msg)



